{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# InferencePipeline on RTSP Stream\n",
    "\n",
    "---\n",
    "[![Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://github.com/roboflow/notebooks/blob/main/notebooks/inferencePipeline.ipynb)\n",
    "\n",
    "The Roboflow Inference Pipeline is a drop-in replacement for the Hosted Inference API that can be deployed on your own hardware. The Inference Pipeline interface is made for streaming and is likely the best route to go for real time use cases. It is an asynchronous interface that can consume many different video sources including local devices (like webcams), RTSP video streams, video files, etc. With this interface, you define the source of a video stream and sinks.\n",
    "\n",
    "We have optimized Inference Pipeline to get maximum performance from the NVIDIA Jetson line of edge-AI devices. We have done this by specifically tailoring the drivers, libraries, and binaries specifically to its CPU and GPU architectures.\n",
    "\n",
    "\n",
    "**Let's begin!**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Install required packages\n",
    "\n",
    "In this cookbook, we'll leverage two Python packages - `inference` and `supervision`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install inference supervision==0.18.0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from inference.core.interfaces.stream.inference_pipeline import InferencePipeline\n",
    "from inference.core.interfaces.stream.sinks import render_boxes\n",
    "import supervision as sv\n",
    "import pandas as pd\n",
    "from collections import defaultdict\n",
    "import cv2\n",
    "import numpy as np\n",
    "import time\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Run Inference Pipeline with COCO Model Aliases & Native FPS Monitor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create an instance of FPSMonitor\n",
    "fps_monitor = sv.FPSMonitor()\n",
    "\n",
    "REGISTERED_ALIASES = {\n",
    "    \"yolov8n-640\": \"coco/3\",\n",
    "    \"yolov8n-1280\": \"coco/9\",\n",
    "    \"yolov8m-640\": \"coco/8\"\n",
    "}\n",
    "\n",
    "API_KEY = \"API_KEY\"\n",
    "RTSP_STREAM = \"RTSP_URL\"\n",
    "\n",
    "# Example alias\n",
    "alias = \"yolov8n-640\"\n",
    "\n",
    "# Function to resolve an alias to the actual model ID\n",
    "def resolve_roboflow_model_alias(model_id: str) -> str:\n",
    "    return REGISTERED_ALIASES.get(model_id, model_id)\n",
    "\n",
    "# Resolve the alias to get the actual model ID\n",
    "model_name = resolve_roboflow_model_alias(alias)\n",
    "\n",
    "# Modify the render_boxes function to enable displaying statistics\n",
    "def on_prediction(predictions, video_frame):\n",
    "    render_boxes(\n",
    "        predictions=predictions,\n",
    "        video_frame=video_frame,\n",
    "        fps_monitor=fps_monitor,  # Pass the FPS monitor object\n",
    "        display_statistics=True,   # Enable displaying statistics\n",
    "    )\n",
    "    \n",
    "pipeline = InferencePipeline.init(\n",
    "    model_id= model_name,\n",
    "    video_reference=RTSP_STREAM,\n",
    "    on_prediction=on_prediction,\n",
    "    api_key=API_KEY,\n",
    "    confidence=0.5,\n",
    ")\n",
    "\n",
    "pipeline.start()\n",
    "pipeline.join()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![InferencePipeline](https://storage.googleapis.com/com-roboflow-marketing/inferencePipeline.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Time in Zone with Bytetrack using Supervision, save data to CSV"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#ByteTrack & Supervision\n",
    "tracker = sv.ByteTrack()\n",
    "annotator = sv.BoxAnnotator()\n",
    "frame_count = defaultdict(int)\n",
    "colors = sv.ColorPalette.default()\n",
    "\n",
    "#define polygon zone of interest\n",
    "polygons = [\n",
    "np.array([\n",
    "[390, 543],[1162, 503],[1510, 711],[410, 819],[298, 551],[394, 543]\n",
    "])\n",
    "]\n",
    "\n",
    "#create zones, zone_annotator, and box_annotator based on polygon zone of interest\n",
    "zones = [\n",
    "    sv.PolygonZone(\n",
    "        polygon=polygon,\n",
    "        frame_resolution_wh=[1440,2560],\n",
    "    )\n",
    "    for polygon\n",
    "    in polygons\n",
    "]\n",
    "zone_annotators = [\n",
    "    sv.PolygonZoneAnnotator(\n",
    "        zone=zone,\n",
    "        color=colors.by_idx(index),\n",
    "        thickness=4,\n",
    "        text_thickness=8,\n",
    "        text_scale=4\n",
    "    )\n",
    "    for index, zone\n",
    "    in enumerate(zones)\n",
    "]\n",
    "box_annotators = [\n",
    "    sv.BoxAnnotator(\n",
    "        color=colors.by_idx(index),\n",
    "        thickness=4,\n",
    "        text_thickness=4,\n",
    "        text_scale=2\n",
    "        )\n",
    "    for index\n",
    "    in range(len(polygons))\n",
    "]\n",
    "\n",
    "\n",
    "#columns for csv output\n",
    "columns = ['trackerID', 'class_id', 'frame_count','entry_timestamp','exit_timestamp','time_in_zone']\n",
    "frame_count_df = pd.DataFrame(columns=columns)\n",
    "\n",
    "# Define a dictionary to store the first detection timestamp for each tracker_id\n",
    "first_detection_timestamps = {}\n",
    "last_detection_timestamps = {}\n",
    "\n",
    "def render(predictions: dict, video_frame) -> None:\n",
    "    detections = sv.Detections.from_roboflow(predictions)\n",
    "    detections = tracker.update_with_detections(detections)\n",
    "    \n",
    "    for zone, zone_annotator, box_annotator in zip(zones, zone_annotators, box_annotators):\n",
    "        mask = zone.trigger(detections=detections)\n",
    "        detections_filtered = detections[mask]\n",
    "        \n",
    "        image = box_annotator.annotate(scene=video_frame.image, detections=detections, skip_label=False)\n",
    "        image = zone_annotator.annotate(scene=image)\n",
    "        \n",
    "        for tracker_id, class_id in zip(detections_filtered.tracker_id, detections_filtered.class_id):\n",
    "            frame_count[tracker_id] += 1\n",
    "            \n",
    "            # Check if tracker_id is not in first_detection_timestamps, if not, add the timestamp\n",
    "            if tracker_id not in first_detection_timestamps:\n",
    "                first_detection_timestamps[tracker_id] = time.time()\n",
    "            \n",
    "            last_detection_timestamps[tracker_id] = time.time()\n",
    "            \n",
    "            time_difference = last_detection_timestamps[tracker_id] - first_detection_timestamps[tracker_id]\n",
    "            \n",
    "            # Add data to the DataFrame\n",
    "            frame_count_df.loc[tracker_id] = [tracker_id, class_id, frame_count[tracker_id], first_detection_timestamps[tracker_id],last_detection_timestamps[tracker_id], time_difference]\n",
    "    \n",
    "    frame_count_df.to_csv('demo.csv', index=False)\n",
    "    \n",
    "    cv2.imshow(\"Prediction\", image)\n",
    "    cv2.waitKey(1)\n",
    "    \n",
    "\n",
    "#Initialize & Deploy InferencePipeline\n",
    "pipeline = InferencePipeline.init(\n",
    "    model_id=\"coco/8\",\n",
    "    video_reference=\"RTSP_URL\",\n",
    "    on_prediction=render,\n",
    "    api_key = 'API_KEY',\n",
    "    confidence=0.5,\n",
    ")\n",
    "pipeline.start()\n",
    "pipeline.join()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![InferencePipelineTracking](https://storage.googleapis.com/com-roboflow-marketing/inferencePipeline_tracking.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
